Greptile Summary

This PR adds support for user-defined operation fusions by introducing register_forward_fusion and register_backward_fusion functions to the public API. The implementation includes a significant refactoring of the op fusion system to improve consistency and maintainability.

Key Changes:

• Added register_forward_fusion() and register_backward_fusion() functions in fuser.py to allow users to register custom fusion functions
• Refactored all fusion functions from standalone functions to static methods on the fused operation classes (e.g., fuse_forward_linear_bias_activation() → ForwardLinearBiasActivation.fuse_forward_ops())
• Simplified fusion function signatures to accept list[FusibleOperation] instead of list[tuple[FusibleOperation, list[int]]], removing index bookkeeping complexity
• Fixed indexing bugs in backward fusion operations (backward_linear_add.py, backward_linear_scale.py) where context indices were incorrect
• Ensured consistent operation order between forward and backward passes by having both work with the same basic operation sequence
• Added comprehensive tests for custom basic operations, custom forward fusions, and custom backward fusions

Impact:
The refactoring provides a cleaner, more maintainable API for operation fusion while fixing several indexing bugs. Users can now define and register their own fusion patterns, extending TransformerEngine's optimization capabilities.

Confidence Score: 5/5

• This PR is safe to merge with no critical issues found
• The PR implements a well-designed feature with comprehensive testing. The refactoring fixes existing indexing bugs in backward operations and significantly improves code maintainability by removing complex index bookkeeping. The new API is clean and follows good design patterns. All changes are backwards compatible since the old fusion functions were internal implementation details.
• No files require special attention

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant User
    participant RegisterAPI as register_forward_fusion/<br/>register_backward_fusion
    participant OperationFuser
    participant FusionFunc as Custom Fusion Function
    participant FusedOp as Fused Operation

    User->>RegisterAPI: register custom fusion function
    RegisterAPI->>OperationFuser: append to forward_fusion_functions<br/>or backward_fusion_functions

    Note over User,FusedOp: During model execution

    User->>OperationFuser: __call__(input, ...)
    OperationFuser->>OperationFuser: maybe_fuse_ops()
    
    loop For each fusion function
        OperationFuser->>FusionFunc: func(ops, recipe=recipe)
        FusionFunc->>FusionFunc: scan ops with sliding window
        alt Pattern matches
            FusionFunc->>FusedOp: create fused op
            FusedOp-->>FusionFunc: fused operation
        else No match
            FusionFunc->>FusionFunc: shift window
        end
        FusionFunc-->>OperationFuser: updated ops list
    end

    OperationFuser->>OperationFuser: store _forward_ops and _backward_ops
    
    Note over OperationFuser,FusedOp: Forward pass
    
    loop For each forward op
        OperationFuser->>FusedOp: fuser_forward(ctxs, input, ...)
        FusedOp->>FusedOp: execute fused computation
        FusedOp->>FusedOp: save state to contexts
        FusedOp-->>OperationFuser: output
    end
    
    Note over OperationFuser,FusedOp: Backward pass
    
    loop For each backward op (reversed)
        OperationFuser->>FusedOp: fuser_backward(ctxs, grad_output, ...)
        FusedOp->>FusedOp: restore from contexts
        FusedOp->>FusedOp: execute fused backward
        FusedOp-->>OperationFuser: grad_input, grad_params
    end
    
    OperationFuser-->>User: outputs and gradients